package com.df.gameai.modules.demo;

import com.df.gameai.config.SimpleChatMemery;
import com.df.gameai.modules.aiservice.MyAiAssistant;
import com.df.gameai.modules.tools.GameAiTools;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.service.TokenStream;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

@RestController
public class ChatController {

    @Autowired
    ChatLanguageModel chatLanguageModel;
    @Autowired
    StreamingChatLanguageModel  streamingChatLanguageModel;
    @Autowired
    OpenAiStreamingChatModel openAiStreamingChatModel;
    ChatMemory memory = null;


    @GetMapping("/chat")
    public Flux<String> model(@RequestParam(value = "msg", defaultValue = "Hello") String msg) {
//        if(memory==null){
//            memory = SimpleChatMemery.builder().id("1").msgs(List.of()).build();
//        }
//        memory.add(UserMessage.from(message));
        Flux<String> flux = Flux.create(sink -> {
            StringBuilder sb = new StringBuilder();
            ChatRequest req = ChatRequest.builder()
                            .toolSpecifications(ToolSpecifications.toolSpecificationsFrom(GameAiTools.class))
                                    .messages(UserMessage.from(msg))
                                            .build();
            streamingChatLanguageModel.chat(req, new StreamingChatResponseHandler() {


                @Override
                public void onPartialResponse(String partialResponse) {
                    sink.next(partialResponse);
                    sb.append(partialResponse);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    sink.complete();
//                    memory.add(AiMessage.from(sb.toString()));
                    System.out.println("resp:"+sb.toString());
                }

                @Override
                public void onError(Throwable throwable) {
                    sink.error(throwable);
                    memory.add(AiMessage.from(sb.toString()));
                }
            });
        });
        return flux;
    }
    @Autowired
    private MyAiAssistant myAiAssistant;

    @GetMapping("/chat2")
    public Flux<String> chat2(@RequestParam(value = "msg", defaultValue = "Hello") String msg) {
        if(memory==null){
            memory = SimpleChatMemery.builder().id("1").msgs(new ArrayList<>()).build();
        }
        memory.add(UserMessage.from(msg));
        Flux<String> flux = Flux.create(sink -> {
            StringBuilder sb = new StringBuilder();
            TokenStream chat = myAiAssistant.chat(msg);
            chat.onPartialResponse(s->{
                sink.next(s);
                sb.append(s);
            }).onCompleteResponse(s->{
                sink.complete();
                memory.add(AiMessage.from(sb.toString()));
            }).onToolExecuted(toolExecution -> {
                System.out.println("工具执行："+toolExecution.request());
                System.out.println("工具结果："+toolExecution.result());
            }).onError(e->{
                sink.error(e);
            }).start();
        });

        return flux;
    }

    @GetMapping("/chat3")
    public Flux<String> chat3(@RequestParam(value = "msg", defaultValue = "Hello") String msg) {
        if(memory==null){
            memory = SimpleChatMemery.builder().id("1").msgs(new ArrayList<>()).build();
        }
        memory.add(UserMessage.from(msg));
        Flux<String> flux = Flux.create(sink -> {
            StringBuilder sb = new StringBuilder();
            TokenStream chat = myAiAssistant.chat2(msg);
            chat.onPartialResponse(s->{
                sink.next(s);
                sb.append(s);
            }).onCompleteResponse(s->{
                sink.complete();
                memory.add(AiMessage.from(sb.toString()));
            }).onError(e->{
                sink.error(e);
            }).start();
        });

        return flux;
    }
}
